

import argparse


def parse_args():

    parser = argparse.ArgumentParser(description="Run ARM_TGNCF")
    parser.add_argument('--is_use_cuda', type=bool, default=True, help='is cuda')
    parser.add_argument('--seed', type=int, default=2020)
    parser.add_argument('--u_top_k', type=int, default=5, help='')
    parser.add_argument('--i_top_k', type=int, default=5, help='')
    parser.add_argument('--wtu', type=int, default=1)
    parser.add_argument('--wti', type=int, default=1)
    parser.add_argument('--lr', type=float, default=0.001)
    parser.add_argument('--train', type=float, default=0.8, help='')
    parser.add_argument('--embed_size', type=int, default=64,
                        help='Embedding size')
    parser.add_argument('--layer_size', nargs='?', default='[32]',
                        help='Output sizes of every layer')
    parser.add_argument('--lysz', type=int, default=2)
    parser.add_argument('--gpu_id', type=int, default=0,
                        help='')
    parser.add_argument('--inter_score', type=float, default=3.1,
                        help='')
    parser.add_argument('--test_start_epoch', type=int, default=0,
                        help='')
    parser.add_argument('--info', type=str, default="")
    parser.add_argument('--decay', type=float, default=0.001, help='')
    parser.add_argument('--mode', type=str, default="O")

    parser.add_argument('--epoch', type=int, default=200)

    parser.add_argument('--batch_size', type=int, default=2048)
    parser.add_argument('--UI_type', type=int, default=3,
                        help='')
    parser.add_argument('--dataset', nargs='?', default='ML_100K',
                        help='Choose a dataset from '
                             '{ML_1M,ML_100K,Amazon_fashion,Musical_Instruments,'
                             'Epinions,BookCrossing}')

    parser.add_argument('--ii_method', type=str, default="ARM",
                        help='ARM'
                             'path_length'
                             'similarity')

    parser.add_argument('--ARM_sequential', type=bool, default=True,
                        help='')

    parser.add_argument('--forecast_method', type=str, default='rate',
                        help='')

    parser.add_argument('--is_in_train', type=bool, default=True,
                        help='')

    parser.add_argument('--events_is_ascend_time', type=bool, default=True,
                        help='')

    parser.add_argument('--model_is_rerun', type=bool, default=True,
                        help='')

    parser.add_argument('--del_cold', type=bool, default=False,
                        help='')

    return parser.parse_args()


ARM_args = parse_args()

inter_scores = {
    'ML_100K': 3,
    'ML_1M': 3,
    'Amazon_fashion': 3,
    'Musical_Instruments': 3,
    'Epinions': 3,
    'BookCrossing': 3
}
original_rating_name = {
    'ML_100K': 'u.data',
    'ML_1M': 'ratings.dat',
    'Amazon_fashion': 'Amazon_fashion.csv',
    'Musical_Instruments': 'Musical_Instruments.csv',
    'Epinions': 'Epinions.dat',
    'BookCrossing': 'BX-Book-Ratings.dat'
}

cold_threshold = {
    'ML_100K': [0, 0],
    'ML_1M': [0, 0],
    'Amazon_fashion': [3, 3],
    'Musical_Instruments': [20, 10],
    'Epinions': [20, 15],
    'BookCrossing': [20, 15]
}
